import datetime
import os
from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow.keras import Model

np.set_printoptions(threshold=np.inf)

cifar10 = tf.keras.datasets.cifar10
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = tf.keras.backend.resize_images(x_train[:100], 7, 7, "channels_last", "nearest").numpy()
# y_train = tf.constant(y_train[:100])
y_train = y_train[:100]
print(type(x_train))
print(type(y_train))
print("x_train.shape", x_train.shape)
print("y_train.shape", y_train.shape)


def h_sigmoid(x):
    output = tf.keras.layers.Activation('hard_sigmoid')(x)
    return output


def h_swish(x):
    output = x * h_sigmoid(x)
    return output


def Squeeze_excitation_layer(x):
    inputs = x
    squeeze = inputs.shape[-1] / 2
    excitation = inputs.shape[-1]
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(squeeze)(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.Dense(excitation)(x)
    x = h_sigmoid(x)
    x = tf.keras.layers.Reshape((1, 1, excitation))(x)
    x = inputs * x

    return x


def BottleNeck(inputs, exp_size, out_size, kernel_size, strides, is_se_existing, activation):
    x = tf.keras.layers.Conv2D(filters=exp_size,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    if activation == "HS":
        x = h_swish(x)
    elif activation == "RE":
        x = tf.keras.layers.Activation(tf.nn.relu6)(x)
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=kernel_size,
                                        strides=strides,
                                        padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    if activation == "HS":
        x = h_swish(x)
    elif activation == "RE":
        x = tf.keras.layers.Activation(tf.nn.relu6)(x)
    if is_se_existing:
        x = Squeeze_excitation_layer(x)
    x = tf.keras.layers.Conv2D(filters=out_size,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation(tf.keras.activations.linear)(x)
    if strides == 1 and inputs.shape[-1] == out_size:
        # print("x:", x)
        # print("inputs:", inputs)
        x = tf.keras.layers.add([x, inputs])

    return x


class MobileNetV3Large(Model):
    def __init__(self, classes):
        super(MobileNetV3Large, self).__init__()
        self.classes = classes
        self.c1 = tf.keras.layers.Conv2D(filters=16,
                                         kernel_size=(3, 3),
                                         strides=2,
                                         padding="same")
        self.b1 = tf.keras.layers.BatchNormalization()

        # self.bNeck1 = BottleNeck(exp_size=16, out_size=16, kernel_size=3, strides=1, is_se_existing=False,
        #                          activation="RE")
        # self.bNeck2 = BottleNeck(exp_size=64, out_size=24, kernel_size=3, strides=2, is_se_existing=False,
        #                          activation="RE")
        # self.bNeck3 = BottleNeck(exp_size=72, out_size=24, kernel_size=3, strides=1, is_se_existing=False,
        #                          activation="RE")
        # self.bNeck4 = BottleNeck(exp_size=72, out_size=40, kernel_size=5, strides=2, is_se_existing=True,
        #                          activation="RE")
        # self.bNeck5 = BottleNeck(exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True,
        #                          activation="RE")
        # self.bNeck6 = BottleNeck(exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True,
        #                          activation="RE")
        # self.bNeck7 = BottleNeck(exp_size=240, out_size=80, kernel_size=3, strides=2, is_se_existing=False,
        #                          activation="HS")
        # self.bNeck8 = BottleNeck(exp_size=200, out_size=80, kernel_size=3, strides=1, is_se_existing=False,
        #                          activation="HS")
        # self.bNeck9 = BottleNeck(exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False,
        #                          activation="HS")
        # self.bNeck10 = BottleNeck(exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False,
        #                           activation="HS")
        # self.bNeck11 = BottleNeck(exp_size=480, out_size=112, kernel_size=3, strides=1, is_se_existing=True,
        #                           activation="HS")
        # self.bNeck12 = BottleNeck(exp_size=672, out_size=112, kernel_size=3, strides=1, is_se_existing=True,
        #                           activation="HS")
        # self.bNeck13 = BottleNeck(exp_size=672, out_size=160, kernel_size=5, strides=2, is_se_existing=True,
        #                           activation="HS")
        # self.bNeck14 = BottleNeck(exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True,
        #                           activation="HS")
        # self.bNeck15 = BottleNeck(exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True,
        #                           activation="HS")

        self.c2 = tf.keras.layers.Conv2D(filters=960,
                                         kernel_size=(1, 1),
                                         strides=1,
                                         padding="same")
        self.b2 = tf.keras.layers.BatchNormalization()
        self.p = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), strides=1)
        self.c3 = tf.keras.layers.Conv2D(filters=1280,
                                         kernel_size=(1, 1),
                                         strides=1,
                                         padding="same")
        self.c4 = tf.keras.layers.Conv2D(filters=self.classes,
                                         kernel_size=(1, 1),
                                         strides=1,
                                         padding="same",
                                         activation=tf.keras.activations.softmax)

    def call(self, inputs):
        x = inputs
        print(x)
        x = self.c1(x)
        x = self.b1(x)
        x = h_swish(x)

        x = BottleNeck(x, exp_size=16, out_size=16, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
        x = BottleNeck(x, exp_size=64, out_size=24, kernel_size=3, strides=2, is_se_existing=False, activation="RE")
        x = BottleNeck(x, exp_size=72, out_size=24, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
        x = BottleNeck(x, exp_size=72, out_size=40, kernel_size=5, strides=2, is_se_existing=True, activation="RE")
        x = BottleNeck(x, exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="RE")
        x = BottleNeck(x, exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="RE")
        x = BottleNeck(x, exp_size=240, out_size=80, kernel_size=3, strides=2, is_se_existing=False, activation="HS")
        x = BottleNeck(x, exp_size=200, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
        x = BottleNeck(x, exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
        x = BottleNeck(x, exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
        x = BottleNeck(x, exp_size=480, out_size=112, kernel_size=3, strides=1, is_se_existing=True, activation="HS")
        x = BottleNeck(x, exp_size=672, out_size=112, kernel_size=3, strides=1, is_se_existing=True, activation="HS")
        x = BottleNeck(x, exp_size=672, out_size=160, kernel_size=5, strides=2, is_se_existing=True, activation="HS")
        x = BottleNeck(x, exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
        x = BottleNeck(x, exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True, activation="HS")

        x = self.c2(x)
        x = self.b2(x)
        x = h_swish(x)
        x = self.p(x)
        x = self.c3(x)
        x = h_swish(x)
        x = self.c4(x)

        return x

# inputs = np.zeros((1, 224, 224, 3), np.float32)
# output_large = MobileNetV3Large(10).call(x_train[:30])
# print(output_large)

model = MobileNetV3Large(10)

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
              metrics=['sparse_categorical_accuracy'])

# checkpoint_save_path = "./checkpoint/MobileNetV3Large.ckpt"
# if os.path.exists(checkpoint_save_path + '.index'):
#     print('-------------load the model-----------------')
#     model.load_weights(checkpoint_save_path)
#
# cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
#                                                  save_weights_only=True,
#                                                  save_best_only=True)
#
# history = model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_test, y_test), validation_freq=1,
#                     callbacks=[cp_callback])

with open('./cost_time.txt', 'w', encoding='utf8') as file:
    file.write("start-time:" + datetime.datetime.now().__str__() + "\n")

history = model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_test, y_test), validation_freq=1)

with open('./cost_time.txt', 'a+', encoding='utf8') as file:
    file.write("end-time:" + datetime.datetime.now().__str__())

model.summary()

###############################################    show   ###############################################
# 显示训练集和验证集的acc和loss曲线
# acc = history.history['sparse_categorical_accuracy']
# val_acc = history.history['val_sparse_categorical_accuracy']
# loss = history.history['loss']
# val_loss = history.history['val_loss']
#
# plt.subplot(1, 2, 1)
# plt.plot(acc, label='Training Accuracy')
# plt.plot(val_acc, label='Validation Accuracy')
# plt.title('Training and Validation Accuracy')
# plt.legend()
#
# plt.subplot(1, 2, 2)
# plt.plot(loss, label='Training Loss')
# plt.plot(val_loss, label='Validation Loss')
# plt.title('Training and Validation Loss')
# plt.legend()
# plt.show()
